Skip to content

feat: select CUDA device consistently across inference backends - #306

Closed
be-student wants to merge 2 commits into
nomadkaraoke:mainfrom
be-student:codex/225-cuda-device-selection
Closed

feat: select CUDA device consistently across inference backends#306
be-student wants to merge 2 commits into
nomadkaraoke:mainfrom
be-student:codex/225-cuda-device-selection

Conversation

@be-student

@be-student be-student commented Sep 6, 2026 •
edited by coderabbitai Bot
Loading

Copy link
Copy Markdown

Summary

Add Separator(cuda_device_index=...) and --cuda_device_index so a multi-GPU user can select the same visible device for PyTorch and ONNX Runtime. configure_cuda(..., device_index=...) also supports an explicit index.

Fixes #225.

The constructor argument is appended to preserve existing positional calls. With no index specified, current automatic behavior remains unchanged. Explicit values are validated before initialization; unavailable CUDA and out-of-range indices raise ValueError rather than silently choosing another backend. Indices are relative to CUDA_VISIBLE_DEVICES.

The change also handles two downstream requirements of indexed devices:

  • MDX compares the provider name when ONNX configuration includes device options, avoiding a false fallback warning.
  • CUDA cache cleanup recognizes indexed devices and enters the selected device context before clearing the cache.

The README documents CLI/API usage and validation behavior.

Validation

  • New tests on unchanged code: 10 failed, 1 passed.
  • Focused API/CLI/runtime tests after implementation: 42 passed, 1 skipped.
  • Final pytest tests/unit -q: 566 passed, 4 skipped, 13 warnings on Python 3.12 with poetry install -E cpu.
  • git diff --check: passed.

Tests cover index propagation to both backends, invalid values, range checks, unavailable CUDA, unchanged defaults, positional compatibility, CLI forwarding, MDX provider options/fallback reporting, and cache cleanup for indexed CUDA devices. These device-selection checks mock CUDA availability and ONNX sessions; real multi-GPU inference was not performed.

AI assistance: Codex implemented this change and ran the listed validation.

Summary by CodeRabbit

  • New Features

    • Select a specific CUDA device through the CLI or Python API.
    • Apply the selected device consistently to PyTorch and ONNX Runtime.
    • Validate device indexes and report clear errors for unavailable or out-of-range devices.
    • Document zero-based indexing and CUDA_VISIBLE_DEVICES behavior.

  • Bug Fixes

    • Improved GPU cache clearing for the active CUDA device.
    • Preserved ONNX Runtime provider configuration during device selection.

@coderabbitai

coderabbitai Bot commented Sep 6, 2026 •
edited
Loading

Copy link
Copy Markdown

Review Change StackReview Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: e8f5a78d-1e8e-4819-9420-a11e590551e0

📥 Commits

Reviewing files that changed from the base of the PR and between 2b04211 and c415d86.

📒 Files selected for processing (4)
  • audio_separator/separator/separator.py
  • tests/unit/test_cli.py
  • tests/unit/test_cuda_device_selection.py
  • tests/unit/test_separator_api_compatibility.py
🚧 Files skipped from review as they are similar to previous changes (4)
  • tests/unit/test_separator_api_compatibility.py
  • tests/unit/test_cuda_device_selection.py
  • audio_separator/separator/separator.py
  • tests/unit/test_cli.py

Included review availability: Your plan provides up to 4 included reviews per hour; 1 remains after this review.

Walkthrough

Changes

CUDA device selection

Layer / File(s) Summary
CUDA selection API and configuration
audio_separator/separator/separator.py, tests/unit/test_cuda_device_selection.py		 Separator validates cuda_device_index, applies it to PyTorch and ONNX Runtime, and raises ValueError for invalid, unavailable, or out-of-range devices.
Runtime provider and cache handling
audio_separator/separator/architectures/mdx_separator.py, audio_separator/separator/common_separator.py, tests/unit/test_cuda_device_selection.py		 ONNX provider tuples retain device options, and GPU cache cleanup runs in the selected CUDA device context.
CLI, documentation, and API compatibility
audio_separator/utils/cli.py, README.md, tests/unit/test_cli.py, tests/unit/test_separator_api_compatibility.py		 The CLI and Python documentation expose the new option. CLI defaults and constructor signature tests include it.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant CLI
  participant Separator
  participant PyTorch
  participant ONNXRuntime
  CLI->>Separator: Pass cuda_device_index
  Separator->>Separator: Validate device index
  Separator->>PyTorch: Select cuda:index
  Separator->>ONNXRuntime: Set CUDA provider device_id
Loading

Merge Risk: ⚪ Minimal · up to c415d

No concrete merge-blocking risk remains in the supplied review context.

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and concisely describes the main change: consistent CUDA device selection across PyTorch and ONNX Runtime inference backends.
Linked Issues check ✅ Passed Issue #225 requires an optional CUDA device index for configure_cuda and consistent use by PyTorch and ONNX Runtime. Separator adds cuda_device_index, validates explicit indices before initializ…
Out of Scope Changes check ✅ Passed The changed files remain within the scope of Issue #225. README and CLI changes expose the new option. clear_gpu_cache and MDX provider handling support indexed CUDA execution. The added tests verif…
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 21 functions across 7 files.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

A rabbit picked a CUDA lane
And checked each index twice
PyTorch hopped to cuda:n
ONNX kept its device slice
The cache cleared clean and nice

Comment @coderabbitai help to get the list of available commands.

@be-student be-student closed this Sep 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Feature]: Allow specifying CUDA device in configure_cuda

1 participant

@be-student